// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/core/query_ctx.h>
#include <pollux/common/base/spill_config.h>
#include <pollux/common/base/trace_config.h>
#include <pollux/common/config/config.h>

namespace kumo::pollux::core {
    // static
    std::shared_ptr<QueryCtx> QueryCtx::create(
        melon::Executor *executor,
        QueryConfig &&queryConfig,
        std::unordered_map<std::string, std::shared_ptr<config::ConfigBase> >
        connectorConfigs,
        cache::AsyncDataCache *cache,
        std::shared_ptr<memory::MemoryPool> pool,
        melon::Executor *spillExecutor,
        const std::string &queryId) {
        std::shared_ptr<QueryCtx> queryCtx(new QueryCtx(
            executor,
            std::move(queryConfig),
            std::move(connectorConfigs),
            cache,
            std::move(pool),
            spillExecutor,
            queryId));
        queryCtx->maybeSetReclaimer();
        return queryCtx;
    }

    QueryCtx::QueryCtx(
        melon::Executor *executor,
        QueryConfig &&queryConfig,
        std::unordered_map<std::string, std::shared_ptr<config::ConfigBase> >
        connectorSessionProperties,
        cache::AsyncDataCache *cache,
        std::shared_ptr<memory::MemoryPool> pool,
        melon::Executor *spillExecutor,
        const std::string &queryId)
        : queryId_(queryId),
          executor_(executor),
          spillExecutor_(spillExecutor),
          cache_(cache),
          connectorSessionProperties_(connectorSessionProperties),
          pool_(std::move(pool)),
          queryConfig_{std::move(queryConfig)} {
        initPool(queryId);
    }

    /*static*/
    std::string QueryCtx::generatePoolName(const std::string &queryId) {
        // We attach a monotonically increasing sequence number to ensure the pool
        // name is unique.
        static std::atomic<int64_t> seqNum{0};
        return fmt::format("query.{}.{}", queryId.c_str(), seqNum++);
    }

    void QueryCtx::maybeSetReclaimer() {
        POLLUX_CHECK_NOT_NULL(pool_);
        POLLUX_CHECK(!underArbitration_);
        if (pool_->reclaimer() != nullptr) {
            return;
        }
        pool_->setReclaimer(QueryCtx::MemoryReclaimer::create(this, pool_.get()));
    }

    void QueryCtx::updateSpilledBytesAndCheckLimit(uint64_t bytes) {
        const auto numSpilledBytes = numSpilledBytes_.fetch_add(bytes) + bytes;
        if (queryConfig_.maxSpillBytes() > 0 &&
            numSpilledBytes > queryConfig_.maxSpillBytes()) {
            POLLUX_SPILL_LIMIT_EXCEEDED(fmt::format(
                "Query exceeded per-query local spill limit of {}",
                succinctBytes(queryConfig_.maxSpillBytes())));
        }
    }

    void QueryCtx::updateTracedBytesAndCheckLimit(uint64_t bytes) {
        if (numTracedBytes_.fetch_add(bytes) + bytes >=
            queryConfig_.queryTraceMaxBytes()) {
            POLLUX_TRACE_LIMIT_EXCEEDED(fmt::format(
                "Query exceeded per-query local trace limit of {}",
                succinctBytes(queryConfig_.queryTraceMaxBytes())));
        }
    }

    std::unique_ptr<memory::MemoryReclaimer> QueryCtx::MemoryReclaimer::create(
        QueryCtx *queryCtx,
        memory::MemoryPool *pool) {
        return std::unique_ptr<memory::MemoryReclaimer>(
            new QueryCtx::MemoryReclaimer(queryCtx->shared_from_this(), pool));
    }

    uint64_t QueryCtx::MemoryReclaimer::reclaim(
        memory::MemoryPool *pool,
        uint64_t targetBytes,
        uint64_t maxWaitMs,
        memory::MemoryReclaimer::Stats &stats) {
        auto queryCtx = ensureQueryCtx();
        if (queryCtx == nullptr) {
            return 0;
        }
        POLLUX_CHECK_EQ(pool->name(), pool_->name());

        const auto leaveGuard =
                melon::makeGuard([&]() { queryCtx->finishArbitration(); });
        queryCtx->startArbitration();
        return memory::MemoryReclaimer::reclaim(pool, targetBytes, maxWaitMs, stats);
    }

    bool QueryCtx::checkUnderArbitration(ContinueFuture *future) {
        POLLUX_CHECK_NOT_NULL(future);
        std::lock_guard<std::mutex> l(mutex_);
        if (!underArbitration_) {
            POLLUX_CHECK(arbitrationPromises_.empty());
            return false;
        }
        arbitrationPromises_.emplace_back("QueryCtx::waitArbitration");
        *future = arbitrationPromises_.back().getSemiFuture();
        return true;
    }

    void QueryCtx::startArbitration() {
        std::lock_guard<std::mutex> l(mutex_);
        POLLUX_CHECK(!underArbitration_);
        POLLUX_CHECK(arbitrationPromises_.empty());
        underArbitration_ = true;
    }

    void QueryCtx::finishArbitration() {
        std::vector<ContinuePromise> promises; {
            std::lock_guard<std::mutex> l(mutex_);
            POLLUX_CHECK(underArbitration_);
            underArbitration_ = false;
            promises.swap(arbitrationPromises_);
        }
        for (auto &promise: promises) {
            promise.setValue();
        }
    }
} // namespace kumo::pollux::core
